import numpy as np
import matplotlib.pyplot as plt
from photutils.psf import EPSFModel
from astropy.stats import sigma_clipped_stats
from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from pathlib import Path
import subprocess

import matplotlib.pyplot as plt
from astropy.io import fits
from pathlib import Path
class SExtractorRunner:
    """
    Run SExtractor and save catalog under result/, check images under result/check/.
    """
    def __init__(self, sex_path='sex', config_path='.', output_path='result',
                 sexconfig=None, sexparamconfig=None):
        self.sex = sex_path
        self.config_path = Path(config_path)
        self.sexconfig = sexconfig or str(self.config_path/'default.sex')
        self.sexparam = sexparamconfig or str(self.config_path/'default.param')
        self.out_dir = Path(output_path)
        (self.out_dir/'check').mkdir(parents=True, exist_ok=True)

    def run(self, img_path: str, out_cat='stars.cat',
            detect_thresh=1.5, analysis_thresh=1.5, detect_minarea=5,
            check_types=None):
        img = Path(img_path)
        catalog = self.out_dir / out_cat
        check_types = check_types or ['SEGMENTATION', 'FILTERED',
                                      'BACKGROUND', 'BACKGROUND_RMS',
                                      'OBJECTS', 'APERTURES']
        check_paths = []
        for c in check_types:
            fname = f"{img.stem}.{c.lower()}.fits"
            check_paths.append(str(self.out_dir/'check'/fname))
        cmd = [
            self.sex, str(img),
            '-c', self.sexconfig,
            '-PARAMETERS_NAME', self.sexparam,
            '-CATALOG_NAME', str(catalog),
            '-DETECT_THRESH', str(detect_thresh),
            '-ANALYSIS_THRESH', str(analysis_thresh),
            '-DETECT_MINAREA', str(detect_minarea),
            '-CHECKIMAGE_TYPE', ','.join(check_types),
            '-CHECKIMAGE_NAME', ','.join(check_paths)
        ]
        subprocess.run(cmd, check=True)
        return str(catalog), check_paths

class PSFExRunner:
    """
    PSFExRunner encapsulates PSFEx model creation with dynamic kernel sizing and output handling.
    """
    def __init__(self, psfex_path='psfex', config_file=None, snr_low=10, kernel_size_factor=7):
        if not config_file:
            raise ValueError('config_file must be provided')
        self.psfex = psfex_path
        self.config = config_file
        self.snr_low = snr_low
        self.kernel_size_factor = kernel_size_factor

    def run(self, catalog_path: str):
        from astropy.table import Table
        from astropy.stats import sigma_clipped_stats

        catalog = Path(catalog_path)
        psf_name = catalog.stem
        psf_dir = catalog.parent / psf_name
        psf_dir.mkdir(exist_ok=True)

        # Read LDAC catalog (extension 2)
        tbl = Table.read(str(catalog), format='fits', hdu=2)
        if 'SNR_WIN' not in tbl.colnames or 'FWHM_IMAGE' not in tbl.colnames:
            raise ValueError("Catalog missing required fields (SNR_WIN, FWHM_IMAGE)")
        sel = (tbl['SNR_WIN'] > self.snr_low) & (tbl['FLAGS'] == 0)
        tbl_sel = tbl[sel]
        if len(tbl_sel) == 0:
            raise RuntimeError("No suitable stars found for PSFEx (check SNR or FLAGS)")

        _, med_fwhm, _ = sigma_clipped_stats(tbl_sel['FWHM_IMAGE'], sigma=3)
        ksize = max(int(med_fwhm) * self.kernel_size_factor, 17)  # Prevent tiny PSF

        # Construct command
        check_types = ['CHI', 'PROTOTYPES', 'SAMPLES', 'RESIDUALS', 'SNAPSHOTS']
        check_names = [str(psf_dir/n) for n in ['chi.fits','proto.fits','samp.fits','resi.fits','snap.fits']]
        checkplots = [str(psf_dir/n) for n in ['fwhm.fits','ellipticity.fits','counts.fits','chi2.fits','resi.fits']]

        cmd = [
            self.psfex, str(catalog), '-c', self.config,
            '-PSF_SIZE', f'{ksize},{ksize}',
            '-PSF_DIR', str(psf_dir),
            '-OUTCAT_NAME', str(psf_dir/'psfex_out.ldac'),
            '-XML_NAME', str(psf_dir/'psfex.xml'),
            '-CHECKPLOT_TYPE', 'FWHM,ELLIPTICITY,COUNTS,CHI2,RESIDUALS',
            '-CHECKPLOT_NAME', ','.join(checkplots),
            '-CHECKIMAGE_TYPE', ','.join(check_types),
            '-CHECKIMAGE_NAME', ','.join(check_names)
        ]
        print('Running PSFEx with:', ' '.join(cmd))
        subprocess.run(cmd, check=True)

        model_path = psf_dir / f'{psf_name}.psf'
        if not model_path.exists():
            raise FileNotFoundError(f'PSF model not found: {model_path}')
        return str(model_path), checkplots

